{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 256, 10, 10])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "\n",
    "class Resnet(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, dim_in, dim_out):\n",
    "        super().__init__()\n",
    "\n",
    "        self.s = torch.nn.Sequential(\n",
    "            torch.nn.GroupNorm(num_groups=32,\n",
    "                               num_channels=dim_in,\n",
    "                               eps=1e-6,\n",
    "                               affine=True),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(dim_in,\n",
    "                            dim_out,\n",
    "                            kernel_size=3,\n",
    "                            stride=1,\n",
    "                            padding=1),\n",
    "            torch.nn.GroupNorm(num_groups=32,\n",
    "                               num_channels=dim_out,\n",
    "                               eps=1e-6,\n",
    "                               affine=True),\n",
    "            torch.nn.SiLU(),\n",
    "            torch.nn.Conv2d(dim_out,\n",
    "                            dim_out,\n",
    "                            kernel_size=3,\n",
    "                            stride=1,\n",
    "                            padding=1),\n",
    "        )\n",
    "\n",
    "        self.res = None\n",
    "        if dim_in != dim_out:\n",
    "            self.res = torch.nn.Conv2d(dim_in,\n",
    "                                       dim_out,\n",
    "                                       kernel_size=1,\n",
    "                                       stride=1,\n",
    "                                       padding=0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #x -> [1, 128, 10, 10]\n",
    "\n",
    "        res = x\n",
    "        if self.res:\n",
    "            #[1, 128, 10, 10] -> [1, 256, 10, 10]\n",
    "            res = self.res(x)\n",
    "\n",
    "        #[1, 128, 10, 10] -> [1, 256, 10, 10]\n",
    "        return res + self.s(x)\n",
    "\n",
    "\n",
    "Resnet(128, 256)(torch.randn(1, 128, 10, 10)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 512, 64, 64])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Atten(torch.nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.norm = torch.nn.GroupNorm(num_channels=512,\n",
    "                                       num_groups=32,\n",
    "                                       eps=1e-6,\n",
    "                                       affine=True)\n",
    "\n",
    "        self.q = torch.nn.Linear(512, 512)\n",
    "        self.k = torch.nn.Linear(512, 512)\n",
    "        self.v = torch.nn.Linear(512, 512)\n",
    "        self.out = torch.nn.Linear(512, 512)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #x -> [1, 512, 64, 64]\n",
    "        res = x\n",
    "\n",
    "        #norm,维度不变\n",
    "        #[1, 512, 64, 64]\n",
    "        x = self.norm(x)\n",
    "\n",
    "        #[1, 512, 64, 64] -> [1, 512, 4096] -> [1, 4096, 512]\n",
    "        x = x.flatten(start_dim=2).transpose(1, 2)\n",
    "\n",
    "        #线性运算,维度不变\n",
    "        #[1, 4096, 512]\n",
    "        q = self.q(x)\n",
    "        k = self.k(x)\n",
    "        v = self.v(x)\n",
    "\n",
    "        #[1, 4096, 512] -> [1, 512, 4096]\n",
    "        k = k.transpose(1, 2)\n",
    "\n",
    "        #[1, 4096, 512] * [1, 512, 4096] -> [1, 4096, 4096]\n",
    "        #0.044194173824159216 = 1 / 512**0.5\n",
    "        #atten = q.bmm(k) * 0.044194173824159216\n",
    "\n",
    "        #照理来说应该是等价的,但是却有很小的误差\n",
    "        atten = torch.baddbmm(torch.empty(1, 4096, 4096, device=q.device),\n",
    "                              q,\n",
    "                              k,\n",
    "                              beta=0,\n",
    "                              alpha=0.044194173824159216)\n",
    "\n",
    "        atten = torch.softmax(atten, dim=2)\n",
    "\n",
    "        #[1, 4096, 4096] * [1, 4096, 512] -> [1, 4096, 512]\n",
    "        atten = atten.bmm(v)\n",
    "\n",
    "        #线性运算,维度不变\n",
    "        #[1, 4096, 512]\n",
    "        atten = self.out(atten)\n",
    "\n",
    "        #[1, 4096, 512] -> [1, 512, 4096] -> [1, 512, 64, 64]\n",
    "        atten = atten.transpose(1, 2).reshape(-1, 512, 64, 64)\n",
    "\n",
    "        #残差连接,维度不变\n",
    "        #[1, 512, 64, 64]\n",
    "        atten = atten + res\n",
    "\n",
    "        return atten\n",
    "\n",
    "\n",
    "Atten()(torch.randn(1, 512, 64, 64)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [0., 0., 0., 0., 0., 0.]],\n",
       "\n",
       "         [[1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [1., 1., 1., 1., 1., 0.],\n",
       "          [0., 0., 0., 0., 0., 0.]]]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Pad(torch.nn.Module):\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.nn.functional.pad(x, (0, 1, 0, 1),\n",
    "                                       mode='constant',\n",
    "                                       value=0)\n",
    "\n",
    "\n",
    "Pad()(torch.ones(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 512, 512])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class VAE(torch.nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = torch.nn.Sequential(\n",
    "            #in\n",
    "            torch.nn.Conv2d(3, 128, kernel_size=3, stride=1, padding=1),\n",
    "\n",
    "            #down\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(128, 128),\n",
    "                Resnet(128, 128),\n",
    "                torch.nn.Sequential(\n",
    "                    Pad(),\n",
    "                    torch.nn.Conv2d(128, 128, 3, stride=2, padding=0),\n",
    "                ),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(128, 256),\n",
    "                Resnet(256, 256),\n",
    "                torch.nn.Sequential(\n",
    "                    Pad(),\n",
    "                    torch.nn.Conv2d(256, 256, 3, stride=2, padding=0),\n",
    "                ),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(256, 512),\n",
    "                Resnet(512, 512),\n",
    "                torch.nn.Sequential(\n",
    "                    Pad(),\n",
    "                    torch.nn.Conv2d(512, 512, 3, stride=2, padding=0),\n",
    "                ),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(512, 512),\n",
    "                Resnet(512, 512),\n",
    "            ),\n",
    "\n",
    "            #mid\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(512, 512),\n",
    "                Atten(),\n",
    "                Resnet(512, 512),\n",
    "            ),\n",
    "\n",
    "            #out\n",
    "            torch.nn.Sequential(\n",
    "                torch.nn.GroupNorm(num_channels=512, num_groups=32, eps=1e-6),\n",
    "                torch.nn.SiLU(),\n",
    "                torch.nn.Conv2d(512, 8, 3, padding=1),\n",
    "            ),\n",
    "\n",
    "            #正态分布层\n",
    "            torch.nn.Conv2d(8, 8, 1),\n",
    "        )\n",
    "\n",
    "        self.decoder = torch.nn.Sequential(\n",
    "            #正态分布层\n",
    "            torch.nn.Conv2d(4, 4, 1),\n",
    "\n",
    "            #in\n",
    "            torch.nn.Conv2d(4, 512, kernel_size=3, stride=1, padding=1),\n",
    "\n",
    "            #middle\n",
    "            torch.nn.Sequential(Resnet(512, 512), Atten(), Resnet(512, 512)),\n",
    "\n",
    "            #up\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(512, 512),\n",
    "                Resnet(512, 512),\n",
    "                Resnet(512, 512),\n",
    "                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),\n",
    "                torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(512, 512),\n",
    "                Resnet(512, 512),\n",
    "                Resnet(512, 512),\n",
    "                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),\n",
    "                torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(512, 256),\n",
    "                Resnet(256, 256),\n",
    "                Resnet(256, 256),\n",
    "                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),\n",
    "                torch.nn.Conv2d(256, 256, kernel_size=3, padding=1),\n",
    "            ),\n",
    "            torch.nn.Sequential(\n",
    "                Resnet(256, 128),\n",
    "                Resnet(128, 128),\n",
    "                Resnet(128, 128),\n",
    "            ),\n",
    "\n",
    "            #out\n",
    "            torch.nn.Sequential(\n",
    "                torch.nn.GroupNorm(num_channels=128, num_groups=32, eps=1e-6),\n",
    "                torch.nn.SiLU(),\n",
    "                torch.nn.Conv2d(128, 3, 3, padding=1),\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    def sample(self, h):\n",
    "        #h -> [1, 8, 64, 64]\n",
    "\n",
    "        #[1, 4, 64, 64]\n",
    "        mean = h[:, :4]\n",
    "        logvar = h[:, 4:]\n",
    "        std = logvar.exp()**0.5\n",
    "\n",
    "        #[1, 4, 64, 64]\n",
    "        h = torch.randn(mean.shape, device=mean.device)\n",
    "        h = mean + std * h\n",
    "\n",
    "        return h\n",
    "\n",
    "    def forward(self, x):\n",
    "        #x -> [1, 3, 512, 512]\n",
    "\n",
    "        #[1, 3, 512, 512] -> [1, 8, 64, 64]\n",
    "        h = self.encoder(x)\n",
    "\n",
    "        #[1, 8, 64, 64] -> [1, 4, 64, 64]\n",
    "        h = self.sample(h)\n",
    "\n",
    "        #[1, 4, 64, 64] -> [1, 3, 512, 512]\n",
    "        h = self.decoder(h)\n",
    "\n",
    "        return h\n",
    "\n",
    "\n",
    "VAE()(torch.randn(1, 3, 512, 512)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'scaling_factor': 0.18215} were passed to AutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import AutoencoderKL\n",
    "\n",
    "#加载预训练模型的参数\n",
    "params = AutoencoderKL.from_pretrained(\n",
    "    'models/diffsion_from_scratch.params', subfolder='vae')\n",
    "\n",
    "vae = VAE()\n",
    "\n",
    "\n",
    "def load_res(model, param):\n",
    "    model.s[0].load_state_dict(param.norm1.state_dict())\n",
    "    model.s[2].load_state_dict(param.conv1.state_dict())\n",
    "    model.s[3].load_state_dict(param.norm2.state_dict())\n",
    "    model.s[5].load_state_dict(param.conv2.state_dict())\n",
    "\n",
    "    if isinstance(model.res, torch.nn.Module):\n",
    "        model.res.load_state_dict(param.conv_shortcut.state_dict())\n",
    "\n",
    "\n",
    "def load_atten(model, param):\n",
    "    model.norm.load_state_dict(param.group_norm.state_dict())\n",
    "    model.q.load_state_dict(param.query.state_dict())\n",
    "    model.k.load_state_dict(param.key.state_dict())\n",
    "    model.v.load_state_dict(param.value.state_dict())\n",
    "    model.out.load_state_dict(param.proj_attn.state_dict())\n",
    "\n",
    "\n",
    "#encoder.in\n",
    "vae.encoder[0].load_state_dict(params.encoder.conv_in.state_dict())\n",
    "\n",
    "#encoder.down\n",
    "for i in range(4):\n",
    "    load_res(vae.encoder[i + 1][0], params.encoder.down_blocks[i].resnets[0])\n",
    "    load_res(vae.encoder[i + 1][1], params.encoder.down_blocks[i].resnets[1])\n",
    "\n",
    "    if i != 3:\n",
    "        vae.encoder[i + 1][2][1].load_state_dict(\n",
    "            params.encoder.down_blocks[i].downsamplers[0].conv.state_dict())\n",
    "\n",
    "#encoder.mid\n",
    "load_res(vae.encoder[5][0], params.encoder.mid_block.resnets[0])\n",
    "load_res(vae.encoder[5][2], params.encoder.mid_block.resnets[1])\n",
    "load_atten(vae.encoder[5][1], params.encoder.mid_block.attentions[0])\n",
    "\n",
    "#encoder.out\n",
    "vae.encoder[6][0].load_state_dict(params.encoder.conv_norm_out.state_dict())\n",
    "vae.encoder[6][2].load_state_dict(params.encoder.conv_out.state_dict())\n",
    "\n",
    "#encoder.正态分布层\n",
    "vae.encoder[7].load_state_dict(params.quant_conv.state_dict())\n",
    "\n",
    "#decoder.正态分布层\n",
    "vae.decoder[0].load_state_dict(params.post_quant_conv.state_dict())\n",
    "\n",
    "#decoder.in\n",
    "vae.decoder[1].load_state_dict(params.decoder.conv_in.state_dict())\n",
    "\n",
    "#decoder.mid\n",
    "load_res(vae.decoder[2][0], params.decoder.mid_block.resnets[0])\n",
    "load_res(vae.decoder[2][2], params.decoder.mid_block.resnets[1])\n",
    "load_atten(vae.decoder[2][1], params.decoder.mid_block.attentions[0])\n",
    "\n",
    "#decoder.up\n",
    "for i in range(4):\n",
    "    load_res(vae.decoder[i + 3][0], params.decoder.up_blocks[i].resnets[0])\n",
    "    load_res(vae.decoder[i + 3][1], params.decoder.up_blocks[i].resnets[1])\n",
    "    load_res(vae.decoder[i + 3][2], params.decoder.up_blocks[i].resnets[2])\n",
    "\n",
    "    if i != 3:\n",
    "        vae.decoder[i + 3][4].load_state_dict(\n",
    "            params.decoder.up_blocks[i].upsamplers[0].conv.state_dict())\n",
    "\n",
    "#decoder.out\n",
    "vae.decoder[7][0].load_state_dict(params.decoder.conv_norm_out.state_dict())\n",
    "vae.decoder[7][2].load_state_dict(params.decoder.conv_out.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.randn(1, 3, 512, 512)\n",
    "\n",
    "a = vae.encoder(data)\n",
    "b = params.encode(data).latent_dist.parameters\n",
    "\n",
    "(a == b).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(True)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.randn(1, 4, 64, 64)\n",
    "\n",
    "a = vae.decoder(data)\n",
    "b = params.decode(data).sample\n",
    "\n",
    "(a == b).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simple",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
